program Adao

	syntax varlist(numeric) [if] , WMatrix(string) [TSls(varname)]

	// Implement [if]
	marksample touse
	preserve
	quietly keep if `touse'
	
	//
	gettoken depvar indeps : varlist
	gettoken bartik controls : indeps
	

display " "
display as text "***********************************************"
display " "
display as text "Calculating standard errors and confidence intervals"
display as text "as from Adao, Kolesar, and Morales (2019)."
display " "
display as text "***********************************************"
display " "
disp as text "Computing step:"
disp as text "5.1"
	
/********************/
/* Remark 5: Step 1 */
/********************/

* Obtain estimates from regressing outcomes onto Bartik and controls.
qui reg `depvar' `bartik' `controls'
* Used in Remark 6.4 to calculate confidence intervals
matrix b_hat = _b[`bartik']

* Obtain epsilon-hat (e1), the estimated regression residuals.
predict e1, residual
mkmat e1, mat(e1) // converts epsilon-hat to a matrix


disp as text "6.1"

/********************/
/* Remark 6: Step 1 */
/********************/

* First part obtained in Remark 5: Step 1, above.

* Regress outcomes minus Bartik (X)*beta0 on controls:.
* Note: beta_0 = 0
qui reg `depvar' `controls'

* Obtain epsilon-hat_b0 (e1_0), the estimated regression residuals.
predict e1_0, residual
mkmat e1_0, mat(e1_0) // converts to a matrix



disp as text "5.2 and 6.2"

/**************************/
/* Remark 5 and 6: Step 2 */
/**************************/

* Regress the Bartik instrument on controls.
qui reg `bartik' `controls'

* Obtain Xumlaut, the estimated regression residuals.
predict Xumlaut, residual
mkmat Xumlaut, mat(Xumlaut) // converts to a matrix

* Defines weight matrix W (used below)
mkmat `wmatrix', mat(W)

* Obtain parameters (Xi-hat) from regressing Xumlaut on W (the weight share).
matrix Xihat = inv(W'*W)*W'*Xumlaut //Equation (25)



disp as text "5.3"

/********************/
/* Remark 5: Step 3 */
/********************/

* We now have epsilon-hat (e1), Xumlaut, and Xihat.
* Next, plug these estimates into equation (26).

* Expression for R: weighted average of residuals by industry_shares for se(beta0) and se(beta)
matrix R = W'*e1 //From Equation (26)

* Number of sector (industry) shares (=#industries * #years) (needed for element-by-element operation below)
local S = rowsof(Xihat)

* Compute element-by-element squaring of Xihat and R.
matrix Xihat2 = Xihat
matrix R2 = R
forval s = 1/`S'{
	matrix Xihat2[`s',1] = (Xihat[`s',1])^2
	matrix R2[`s',1] = (R[`s',1])^2
}

* Effectively generates element-by-element square of Xumlaut
gen Xumlaut2 = Xumlaut^2
mkmat Xumlaut2, mat(Xumlaut2)
				
* Numerator of se(beta-hat) in Equation (26)
matrix Num2 = R2'*Xihat2
* Denominator of se(beta-hat) in Equation (26) and (27)
mata : st_matrix("Den", sum(st_matrix("Xumlaut2")))
* Standard Errors (se) from Equation (26)
matrix se = sqrt(Num2[1,1])/Den[1,1]
local se : di %5.4f se[1,1]


disp as text "6.3"

/********************/
/* Remark 6: Step 3 */
/********************/

* We now have epsilon-hat_beta0 (e1_0), Xumlaut, and Xihat.
* Next, plug these estimates into equation (26).

* Expression for R: weighted average of residuals by industry_shares for se(beta0) and se(beta)
matrix R0 = W'*e1_0

* Compute element-by-element squaring of R0.
matrix R02 = R0
forval s = 1/`S'{
	matrix R02[`s',1] = (R0[`s',1])^2
}
				
* Numerator of se(beta0-hat) in Equation (27)
matrix Num02 = R02'*Xihat2
* Standard Errors (se0) from Equation (27)
matrix se0 = sqrt(Num02[1,1])/Den[1,1]
local se0 : di %5.4f se0[1,1]


disp as text "6.4"

/********************/
/* Remark 6: Step 4 */
/********************/
	
* First component of Q
matrix Qa = (Xumlaut'*Xumlaut)*(Xumlaut'*Xumlaut)/1.96^2
	
* Second component of Q
matrix WX = W'*Xumlaut //Inner product of weight shares and Xumlaut
* Compute element-by-element squaring of R0.
matrix WX2 = WX
forval s = 1/`S'{
	 matrix WX2[`s',1]= (WX2[`s',1])^2
}
matrix Qb = Xihat2'*WX2

* Combining components of Q
matrix Q = Qa - Qb

* Build scalar A
matrix A = 0
forval s = 1/`S'{
	matrix A[1,1]= A[1,1] + Xihat2[`s',1]*R[`s',1]*WX[`s',1]/Q[1,1]
}

* Building the term of 9Xumlaut*Xumlaut)^2 in the denominator of the denominator of the square room term.
matrix XX2 = (Xumlaut'*Xumlaut)*(Xumlaut'*Xumlaut)

matrix ciL = b_hat[1,1] - A[1,1] - sqrt(A[1,1]^2 + (se[1,1]^2)/(Q[1,1]/XX2[1,1]))
local ciL : di %5.4f ciL[1,1]
matrix ciH = b_hat[1,1] - A[1,1] + sqrt(A[1,1]^2 + (se[1,1]^2)/(Q[1,1]/XX2[1,1]))
local ciH : di %5.4f ciH[1,1]


display " "
display as text "***********************************************"
display " "
display as result "The OLS standard error is `se'."
display " "
display as result "The OLS standard error under the null hypothesis is `se0'."
display " "
display as result "The OLS confidence interval is [`ciL',`ciH']."
display " "

	
/* OLS or 2SLS */
* If a no TSls variable is provided, code will end at this step.
if missing("`tsls'") {
	display "Using OLS only"
}


display as text "***********************************************"
display " "
disp as text "2SLS Standard Errors and Confidence Intervals"

/*************************************************/
/* 2SLS Standard Errors and Confidence Intervals */
/*************************************************/

else {
	display as text "Using OLS and 2SLS."
	
	*First Stage Coefficients and Residuals
	qui reg `tsls' `bartik' `controls'
	matrix b_fs = _b[`bartik']
	predict efs, residual
	
	* Used to calculate the denominator in equation (33)
	mkmat `tsls', matrix(Y2)
	mkmat `depvar', matrix(Y1)
	matrix XY2 = Xumlaut'*Y2
	matrix XY1 = Xumlaut'*Y1

	*Two Stage Estimate: Eq 33
	matrix alphahat = XY1[1,1]/XY2[1,1]

	* Edelta equal to the difference in the BartikOLS residuals and first stage
	mkmat `depvar', matrix(Y1)

	* Used to calculate ehat,
	* where ehat = Y1 - Y2*alphahat - Z'(Z'Z)^(-1)Z'(Y1 - Y2*alphahat)
	matrix edelta = Y1 - Y2*alphahat[1,1] 
	svmat edelta, name(edelta)
	
	* Note: Z = `controls'
	* This calculates Z'(Z'Z)^(-1)Z'(Y1 - Y2*alphahat), for use in ehat.
	qui reg edelta `controls'
	predict Ze, xb
	mkmat Ze, matrix(Ze)

	* This completes the calculation of ehat.
	matrix ehat = edelta - Ze
	
	*Rhat terms from Equation 36
	matrix Rhat = W'*ehat
	matrix Rhat2 = Rhat
	* Compute element-by-element squaring of Rhat.
	forval s = 1/`S'{
		matrix Rhat2[`s',1]= (Rhat[`s',1])^2
	}

	
	*Expression for se(alpha) (see equation (36))
	matrix se_alpha = sqrt(Num2[1,1])/abs(XY2[1,1])
	local se_alpha : di %5.4f se_alpha[1,1]

	matrix ciL2SLS = alphahat[1,1] - A[1,1] - sqrt((A[1,1]^2 + se_alpha[1,1]^2 / (Q[1,1]/XX2[1,1])))
	local ciL2SLS : di %5.4f ciL2SLS[1,1]
	matrix ciH2SLS = alphahat[1,1] - A[1,1] + sqrt((A[1,1]^2 + se_alpha[1,1]^2 / (Q[1,1]/XX2[1,1])))
	local ciH2SLS : di %5.4f ciH2SLS[1,1]
	
	

display " "
display as text "***********************************************"
display " "
display as result "The 2SLS standard error is `se_alpha'."
display " "
display as result "The 2SLS confidence interval is [`ciL2SLS',`ciH2SLS']."
display " "
display as text "***********************************************"
display " "
	
}
	
end